package com.whut.coc.application.config;

import com.whut.coc.domain.repository.RedisChatMemoryStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.data.annotation.Reference;

import java.util.ArrayList;
import java.util.List;

/**
 * @author wangzonghe
 * @date 2025/6/20 17:08
 */
@Configuration
@RequiredArgsConstructor
public class CommonConfig {

    private final OpenAiChatModel model;

    @Reference
    private final RedisChatMemoryStore redisChatMemoryStore;

    private final EmbeddingModel embeddingModel;

    /**
     * 构建会话记忆对象
     * @return 会话记忆对象
     */
    @Bean
    public ChatMemory chatMemory() {
        return MessageWindowChatMemory.builder()
                .maxMessages(20)
                .build();
    }

    /**
     * 构建ChatMemoryProvider对象
     * @return ChatMemoryProvider对象 根据id
     */
    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(20)
                .chatMemoryStore(redisChatMemoryStore)
                .build();
    }

    /**
     * 构建向量库操作对象
     * @return EmbeddingStore对象
     */
    @Bean
    @Primary
    public EmbeddingStore store() {

        // 1.加载文档进内存
        List<Document> documents = new ArrayList<>();
        List<Document> contents = ClassPathDocumentLoader.loadDocuments("document/content");
        List<Document> pdfs = ClassPathDocumentLoader.loadDocuments("document/pdf", new ApachePdfBoxDocumentParser());
        documents.addAll(contents);
        documents.addAll(pdfs);
        // 2. 构建向量库操作对象 (操作的是内存版本的向量数据库)
        InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();
        
        // 构建文档分割器对象
        DocumentSplitter ds = DocumentSplitters.recursive(500, 100);

        // 3. 构建一个EmbeddingStoreIngestor对象，玩名称文本数据切割与向量化存储
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .embeddingStore(store)
                .documentSplitter(ds)
                .embeddingModel(embeddingModel)
                .build();
        ingestor.ingest(documents);

        return store;

    }

    /**
     * 构建内容检索对象
     * @return ContentRetriever对象
     */
    @Bean
    public ContentRetriever contentRetriever(EmbeddingStore store) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(store)
                .minScore(0.5)
                .maxResults(3)
                .embeddingModel(embeddingModel)
                .build();
    }

}
